use std::{fmt::Display, num::ParseFloatError};

use lalrpop_util::ParseError;
use logos::{Logos, SpannedIter};

use crate::{chem_parse, parse::element::TABLE};

#[derive(Logos, Debug, PartialEq)]
#[logos(error = LexicalError)]
#[logos(skip  r"[ \t]+")]
pub enum ExprToken {
    #[regex(r"((\d*\.\d+)|(\d+\.?))", |input| input.slice().parse::<f64>())]
    Number(f64),
    #[regex(r"(<[^>]*>)|(#[^#]*#)|('[^']*')", |input| chem_parse(input.slice().trim_matches(['<', '>', '#', '\''])))]
    Chemical(f64),
    #[token("+")]
    Add,
    #[token("-")]
    Sub,
    #[token("*")]
    Mul,
    #[token("/")]
    Div,
    #[regex(r"[({\[（【]")]
    Lparent,
    #[regex(r"[)}\]）】]")]
    Rparent,
}

impl ExprToken {
    #[inline]
    pub fn precedence(&self) -> u8 {
        match self {
            ExprToken::Add | ExprToken::Sub => 1,
            ExprToken::Mul | ExprToken::Div => 2,
            _ => 0,
        }
    }
    pub fn apply(&self, rhs: f64, lhs: f64) -> f64 {
        match self {
            ExprToken::Add => lhs + rhs,
            ExprToken::Sub => lhs - rhs,
            ExprToken::Mul => lhs * rhs,
            ExprToken::Div => lhs / rhs,
            _ => unreachable!(),
        }
    }
}

#[derive(Logos, Debug, PartialEq, Copy, Clone)]
#[logos(error = LexicalError)]
pub enum ChemToken {
    #[regex(r"((\d*\.\d+)|(\d+\.?))", |input| input.slice().parse::<f64>())]
    Number(f64),
    #[regex(r"[A-Z][a-z]?", |input| {
        TABLE.get_mw(input.slice())
    })]
    Elem(f64),
    #[regex(r"[({\[] *")]
    Lparent,
    #[regex(r" *[)}\]]")]
    Rparent,
    #[regex(r"[ \t]+")]
    Space,
}

impl Display for ChemToken {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Self::Number(num) => write!(f, "{}", num),
            Self::Elem(elem) => write!(f, "{}", elem),
            Self::Lparent => write!(f, "("),
            Self::Rparent => write!(f, ")"),
            Self::Space => write!(f, "Space"),
        }
    }
}

#[derive(Debug, Default, Clone, PartialEq)]
pub enum LexicalError {
    InvalidFloat(ParseFloatError),
    #[default]
    InvalidToken,
    InvalidElement(String),
}

impl From<ParseFloatError> for LexicalError {
    fn from(err: ParseFloatError) -> Self {
        LexicalError::InvalidFloat(err)
    }
}

impl Into<LexicalError> for ParseError<usize, ChemToken, LexicalError> {
    fn into(self) -> LexicalError {
        match self {
            Self::User { error } => error,
            _ => LexicalError::InvalidToken,
        }
    }
}

impl Display for LexicalError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidElement(element) => write!(f, "Cannot parse {} into an element", element),
            Self::InvalidFloat(err) => write!(f, "{}", err),
            Self::InvalidToken => write!(f, "Unknow character"),
        }
    }
}

pub struct ChemicalLexer<'input> {
    pub token_stream: SpannedIter<'input, ChemToken>,
}

impl<'input> ChemicalLexer<'input> {
    pub fn new(input: &'input str) -> Self {
        Self {
            token_stream: ChemToken::lexer(input).spanned(),
        }
    }
}

pub type Spanned<Tok, Loc, Error> = Result<(Loc, Tok, Loc), Error>;

impl<'input> Iterator for ChemicalLexer<'input> {
    type Item = Spanned<ChemToken, usize, LexicalError>;
    fn next(&mut self) -> Option<Self::Item> {
        self.token_stream
            .next()
            .map(|(token, span)| Ok((span.start, token?, span.end)))
    }
}
